package com.example.agent_springai.util;

import com.example.agent_springai.entity.po.DrawChartParam;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import kong.unirest.HttpResponse;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;
import java.util.List;
import java.util.ArrayList;
import java.util.LinkedHashMap;

@Slf4j
public class MCPDrawUtil {

    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 构建MCP请求（完全对齐测试类格式）
     */
    public Map<String, Object> buildMCPRequest(DrawChartParam param) {
        Map<String, Object> requestMap = new HashMap<>();
        requestMap.put("jsonrpc", "2.0");
        requestMap.put("id", "req-" + System.currentTimeMillis());
        requestMap.put("method", "tools/call");

        // params节点
        Map<String, Object> params = new HashMap<>();

        // 根据图表类型设置工具名称
        String toolName = getToolNameByType(param.getType());
        params.put("name", toolName);

        // 处理数据中的时间戳，解决同一时间多数据问题
        Object processedData = processDataWithTimestamps(param.getData());

        // arguments节点（核心业务数据）
        Map<String, Object> arguments = new HashMap<>();
        arguments.put("data", processedData);
        arguments.put("theme", param.getTheme());
        arguments.put("style", param.getStyle());
        arguments.put("width", param.getWidth() == null ? 600 : param.getWidth());
        arguments.put("height", param.getHeight() == null ? 400 : param.getHeight());
        arguments.put("title", param.getTitle().trim());

        params.put("arguments", arguments);

        // _meta节点
        Map<String, Object> meta = new HashMap<>();
        meta.put("progressToken", 1);
        params.put("_meta", meta);

        requestMap.put("params", params);
        return requestMap;
    }

    /**
     * 处理数据中的时间戳，解决同一时间多数据问题
     * @param data 原始数据
     * @return 处理后的数据
     */
    private Object processDataWithTimestamps(Object data) {
        try {
            if (data instanceof List) {
                List<Map<String, Object>> dataList = (List<Map<String, Object>>) data;
                
                // 检查是否包含timestamp字段
                boolean hasTimestamp = dataList.size() > 0 && dataList.get(0).containsKey("timestamp");
                
                if (hasTimestamp) {
                    // 使用timestamp字段作为time值
                    List<Map<String, Object>> processedDataList = new ArrayList<>();
                    for (Map<String, Object> item : dataList) {
                        Map<String, Object> newItem = new HashMap<>(item);
                        // 将timestamp转换为适合图表显示的格式
                        if (item.containsKey("timestamp") && item.get("timestamp") != null) {
                            Long timestamp = ((Number) item.get("timestamp")).longValue();
                            // 转换为日期时间格式
                            java.time.Instant instant = java.time.Instant.ofEpochMilli(timestamp);
                            java.time.ZoneId zoneId = java.time.ZoneId.systemDefault();
                            java.time.format.DateTimeFormatter formatter = 
                                java.time.format.DateTimeFormatter.ofPattern("HH:mm:ss.SSS");
                            String formattedTime = java.time.ZonedDateTime.ofInstant(instant, zoneId).format(formatter);
                            newItem.put("time", formattedTime);
                        }
                        processedDataList.add(newItem);
                    }
                    return processedDataList;
                } else {
                    // 原有逻辑：处理重复时间点
                    Map<String, Integer> timeGroupCount = new LinkedHashMap<>();
                    
                    // 统计每个时间点的数据数量
                    for (Map<String, Object> item : dataList) {
                        String time = (String) item.get("time");
                        timeGroupCount.put(time, timeGroupCount.getOrDefault(time, 0) + 1);
                    }
                    
                    // 如果有重复时间点，则添加微小的时间偏移
                    Map<String, Integer> timeOffset = new HashMap<>();
                    List<Map<String, Object>> processedDataList = new ArrayList<>();
                    
                    for (Map<String, Object> item : dataList) {
                        String time = (String) item.get("time");
                        Map<String, Object> newItem = new HashMap<>(item);
                        
                        if (timeGroupCount.get(time) > 1) {
                            int offset = timeOffset.getOrDefault(time, 0);
                            if (offset > 0) {
                                // 为重复的时间点添加毫秒偏移
                                String newTime = time + "." + String.format("%03d", offset * 100);
                                newItem.put("time", newTime);
                            }
                            timeOffset.put(time, offset + 1);
                        }
                        
                        processedDataList.add(newItem);
                    }
                    
                    return processedDataList;
                }
            }
        } catch (Exception e) {
            log.warn("处理时间戳时出错，使用原始数据: {}", e.getMessage());
        }
        
        // 出现异常时返回原始数据
        return data;
    }

    /**
     * 根据图表类型获取MCP工具名称
     */
    public String getToolNameByType(String type) {
        switch (type.toLowerCase()) {
            case "pie":
                return "generate_pie_chart";
            case "bar":
                return "generate_bar_chart";
            case "line":
            default:
                return "generate_line_chart";
        }
    }

    /**
     * 解析MCP响应，提取图片URL
     */
    public String parseMCPResponse(HttpResponse<String> response) {
        if (!response.isSuccess()) {
            log.error("请求失败，状态码: {}, 响应体: {}", response.getStatus(), response.getBody());
            return "画图请求失败，状态码: " + response.getStatus();
        }

        String responseBody = response.getBody();
        log.info("MCP原始响应: {}", responseBody);

        try {
            // 处理SSE响应格式：提取data后的JSON
            String dataJson = extractSseData(responseBody);
            if (dataJson == null) {
                log.error("无法提取SSE数据，响应格式异常");
                return "画图失败：响应格式异常";
            }

            JsonNode jsonResponse = objectMapper.readTree(dataJson);

            if (jsonResponse.has("result")) {
                // 提取图片URL
                String imageUrl = jsonResponse.get("result")
                        .get("content")
                        .get(0)
                        .get("text")
                        .asText();

                log.info("生成图表成功，图片地址: {}", imageUrl);
                // 只返回图片链接
                return imageUrl;

            } else if (jsonResponse.has("error")) {
                String errorMsg = jsonResponse.get("error").get("message").asText();
                log.error("MCP服务报错: {}", errorMsg);
                return "画图失败：" + errorMsg;
            } else {
                log.error("未知响应格式: {}", jsonResponse);
                return "画图失败：未知响应格式";
            }

        } catch (Exception e) {
            log.error("解析MCP响应异常", e);
            return "画图失败：解析响应异常 - " + e.getMessage();
        }
    }

    /**
     * 提取SSE响应中data字段的内容
     */
    public String extractSseData(String sseResponse) {
        if (sseResponse == null) {
            return null;
        }

        // 处理SSE格式：data: {json}
        if (sseResponse.contains("data: ")) {
            String[] parts = sseResponse.split("data: ");
            if (parts.length > 1) {
                return parts[1].trim();
            }
        }

        // 如果不是SSE格式，直接返回原响应
        return sseResponse.trim();
    }
}